
from gym import ObservationWrapper as Environment
from .policy import Agent
import torch


class CartPole(Environment):
    def __init__(self, env=None):
        if env is None:
            import gym
            env = gym.make('CartPole-v0')
        super(CartPole, self).__init__(env)

    def close(self):
        self.env.close()

    def render(self):
        return _get_screen(self.env)

    def observation(self, observation):
        screen = self.render()
        if hasattr(self, 'last_screen'):
            state = screen - self.last_screen
        else:
            state = screen
        self.last_screen = screen
        return state


def evaluate(agent, env, n_episode=100):
    sum = .0
    for i in range(n_episode):
        reward = agent.play(env, verbose=False)
        sum += reward
        print(f'episode {i}: reward={reward}, average={sum / (i + 1)}')
    return sum / n_episode


def _get_cart_location(env, screen_width):
    world_width = env.x_threshold * 2
    scale = screen_width / world_width
    return int(env.state[0] * scale + screen_width / 2.0)


def _get_screen(environment):
    import numpy as np
    screen = environment.render('rgb_array')
    screen = screen.transpose((2, 0, 1))
    _, screen_height, screen_width = screen.shape
    screen = screen[:, int(screen_height*0.4):int(screen_height*0.8)]
    view_width = int(screen_width * 0.6)
    cart_location = _get_cart_location(environment, screen_width)
    if cart_location < view_width // 2:
        slice_range = slice(view_width)
    elif cart_location > (screen_width - view_width // 2):
        slice_range = slice(-view_width, None)
    else:
        slice_range = slice(cart_location - view_width // 2,
                            cart_location + view_width // 2)
    screen = screen[:, :, slice_range]
    screen = np.ascontiguousarray(screen, dtype=np.float32) / 255
    return torch.from_numpy(screen).unsqueeze(0)


if __name__ == '__main__':
    player = Agent()
    env = CartPole()
    average = evaluate(player, env)
    print(f'average = {average}')
